import ctypes
from enum import Enum, Flag, auto
from math import ceil


class RaspModbusRegisterAddressEnum(Enum):
    """
    树莓派的modbus寄存器地址列表，见【关于树莓派modbus通信协议】
    """
    REG1 = 0x0066  # 水位寄存器


class ZsModbusRegisterAddressEnum(Enum):
    """
    众山DTU和RTU的寄存器地址列表，见【众山DTU Modbus协议手册】
    """
    AI1_H = 0x0000
    AI1_L = 0x0001
    AI2_H = 0x0002
    AI2_L = 0x0003
    AI3_H = 0x0004
    AI3_L = 0x0005
    AI4_H = 0x0006
    AI4_L = 0x0007
    AI5_H = 0x0008
    AI5_L = 0x0009
    AI6_H = 0x000A
    AI6_L = 0x000B
    AI7_H = 0x000C
    AI7_L = 0x000D
    AI8_H = 0x000E
    AI8_L = 0x000F
    DI1 = 0x0010
    DI2 = 0x0011
    DI3 = 0x0012
    DI4 = 0x0013
    DO1 = 0x0014
    DO2 = 0x0015
    DO3 = 0x0016
    DO4 = 0x0017


class ZzModbusRegisterAddressEnum(Enum):
    """
    继电器模块的寄存器地址列表，见【ZZ-IO1600-串口版使用说明书】
    """
    DO1 = 0x0000
    DO2 = 0x0001
    DO3 = 0x0002
    DO4 = 0x0003
    DO5 = 0x0004
    DO6 = 0x0005
    DO7 = 0x0006
    DO8 = 0x0007
    DO9 = 0x0008
    DO10 = 0x0009
    DO11 = 0x000A
    DO12 = 0x000B
    DO13 = 0x000C
    DO14 = 0x000D
    DO15 = 0x000E
    DO16 = 0x000F


class ModbusStructureTypeEnum(Enum):
    """
    Modbus协议的字段类型
    """
    DEVICE_ADDRESS = auto()  # 器件地址
    COMMAND = auto()  # 功能码
    REGISTER_START_ADDRESS = auto()  # 寄存器（起始）地址
    CRC_CHECK = auto()  # CRC校验
    COIL_NUMBER = auto()  # 线圈数量
    DI_NUMBER = auto()  # DI数量
    REGISTER_NUMBER = auto()  # 寄存器数量
    COIL_VALUE = auto()  # 线圈值
    REGISTER_VALUE = auto()  # 寄存器值


class ModbusCommandEnum(Enum):
    """
    Modbus功能码
    """
    READ_COIL = 0x01  # 读线圈
    READ_DI = 0x02  # 读DI
    READ_HOLDING_REGISTER = 0x03  # 读保持寄存器
    READ_INPUT_REGISTER = 0x04  # 读输入寄存器
    WRITE_SINGLE_COIL = 0x05  # 控制单个线圈
    WRITE_SINGLE_REGISTER = 0x06  # 写单个寄存器
    WRITE_MULTIPLE_COIL = 0x0F  # 控制多个线圈
    WRITE_MULTIPLE_REGISTER = 0x10  # 写多个寄存器


class ModbusErrorCodeEnum(Enum):
    """
    Modbus响应错误码
    """
    COMMAND_NOT_SUPPORTED = 0x01  # 不支持的功能码
    REGISTER_ADDRESS_OVERFLOW = 0x02  # 寄存器地址错误
    REGISTER_NUMBER_OR_REGISTER_VALUE_ILLEGAL = 0x03  # 寄存器数量或值错误
    REGISTER_READ_WRITE_ERROR = 0x04  # 寄存器读写错误


class ModbusErrorCommandTypeEnum(Enum):
    """
    Modbus响应的错误功能码
    与ModbusCommandEnum对应
    """
    READ_COIL_ERROR = 0x81
    READ_DI_ERROR = 0x82
    READ_HOLDING_REGISTER_ERROR = 0x83
    READ_INPUT_REGISTER_ERROR = 0x84
    WRITE_SINGLE_COIL_ERROR = 0x85
    WRITE_SINGLE_REGISTER_ERROR = 0x86
    WRITE_MULTIPLE_COIL_ERROR = 0x8F
    WRITE_MULTIPLE_REGISTER = 0x90


class ModbusWriteSingleCoilValueEnum(Enum):
    """
    Modbus写单个线圈值时用的状态值
    """
    ON = 0xFF00  # 闭合
    OFF = 0x0000  # 断开


class ModbusWriteSingleRegisterValueEnum(Enum):
    """
    Modbus写单个寄存器时用的状态值
    """
    ON = 0x0001
    OFF = 0x0000


class ModbusWriteMultipleCoilValueEnum(Flag):
    """
    Modbus写多个线圈时用的状态
    线圈1闭合，线圈2断开，线圈3闭合，线圈4断开表达为：
    ModbusWriteMultipleCoilValueEnum.DO1_ON | ModbusWriteMultipleCoilValueEnum.DO3_ON
    """
    DO1_ON = 0x01
    DO2_ON = 0x02
    DO3_ON = 0x04
    DO4_ON = 0x08


class ModbusReadCoilValueEnum(Enum):
    """
    通过Modbus读取线圈时返回的状态
    """
    OFF = 0  # 断开
    ON = 1  # 闭合

    def __str__(self):
        return self.value.__str__()


class ModbusReadDiValueEnum(Enum):
    """
    通过Modbus读取DI时返回的状态
    """
    HIGH = 1  # 高电平
    LOW = 0  # 低电平

    def __str__(self):
        return self.value.__str__()


class ModbusData(object):
    """
    服务器端所需的众山DTU、RTU的Modbus编码和解码逻辑
    符合【众山DTU Modbus协议手册】的描述
    """

    def __init__(self, blob: bytes = None, decoded: dict = None):
        """
        进行Modbus解码和编码
        提供blob则解码，提供decoded则编码
        :param blob: 编码后的Modbus协议bytes
        :param decoded: 解码后的内容，具体格式和示例见modbus.py的main
        """
        self.blob = blob  # 用于输入或存储编码后的Modbus协议bytes
        self.decoded = decoded  # 用于存储输入或解码后的内容
        self.crc16 = ctypes.cdll.LoadLibrary("./libcrc16.so").crc16
        if blob is None and decoded is not None:
            self._encode()
        elif blob is not None and decoded is None:
            self._decode()
        else:
            raise ValueError("blob or decode")

    def get_blob(self):
        return self.blob

    def get_decoded(self):
        return self.decoded

    def _decode(self):
        """
        外部勿调用
        仅实现众山的协议 【众山DTU Modbus协议手册】读线圈和读离散输入响应部分
        :return: 无返回
        """
        self.decoded = {}
        self.bytes_list = list(self.blob)
        if self.crc16(self.blob[:-2], len(self.blob[:-2])) != ((self.bytes_list[-2] << 8) | self.bytes_list[-1]):
            raise AssertionError("CRC verification failed")
        self.decoded[ModbusStructureTypeEnum.DEVICE_ADDRESS] = self.bytes_list[0]
        try:
            command = ModbusCommandEnum(self.bytes_list[1])
        except Exception:
            command = ModbusErrorCommandTypeEnum(self.bytes_list[1])
            raise ConnectionError(f"{command.name}, received: {self.blob.hex(sep=' ')}")

        self.decoded[ModbusStructureTypeEnum.COMMAND] = command
        if command == ModbusCommandEnum.READ_COIL:
            bytes_number = self.bytes_list[2]
            if bytes_number > 8:
                raise AssertionError("READ_COIL bytes_number > 8")
            coil_value = 0
            for i in range(3, 3 + bytes_number):
                coil_value >>= 8
                coil_value |= (self.bytes_list[i] << ((bytes_number - 1) << 3))
            for i in range(bytes_number << 3):
                self.decoded["coil" + str(i)] = ModbusReadCoilValueEnum((coil_value >> i) & 0x01)
        elif command == ModbusCommandEnum.READ_DI:
            bytes_number = self.bytes_list[2]
            if bytes_number != 1:
                raise AssertionError("READ_DI bytes_number != 1")
            di_value = self.bytes_list[3]
            self.decoded["di0"] = ModbusReadDiValueEnum(di_value & 0x01)
            self.decoded["di1"] = ModbusReadDiValueEnum((di_value >> 1) & 0x01)
            self.decoded["di2"] = ModbusReadDiValueEnum((di_value >> 2) & 0x01)
            self.decoded["di3"] = ModbusReadDiValueEnum((di_value >> 3) & 0x01)
        elif command == ModbusCommandEnum.READ_HOLDING_REGISTER:
            raise NotImplementedError("READ_HOLDING_REGISTER not implemented")
        elif command == ModbusCommandEnum.READ_INPUT_REGISTER:
            bytes_number = self.bytes_list[2]
            reg_value = int.from_bytes(bytes(self.bytes_list[3:3 + bytes_number]), byteorder='big', signed=False)
            self.decoded["reg1"] = reg_value
        elif command == ModbusCommandEnum.WRITE_SINGLE_COIL:
            pass
        elif command == ModbusCommandEnum.WRITE_SINGLE_REGISTER:
            pass
        elif command == ModbusCommandEnum.WRITE_MULTIPLE_COIL:
            pass
        elif command == ModbusCommandEnum.WRITE_MULTIPLE_REGISTER:
            pass

    def _encode_stage_1(self):
        """
        外部勿调用
        和self._encode_stage_2()一起，编码形成Modbus协议的bytes
        这个方法读取命令、地址等信息，并把命令存在self._command里供self._encode_stage_2()读取
        :return: 无返回
        """
        self.bytes_list = []
        self._command = self.decoded[ModbusStructureTypeEnum.COMMAND]

        device_address = self.decoded[ModbusStructureTypeEnum.DEVICE_ADDRESS]
        register_start_address = self.decoded[ModbusStructureTypeEnum.REGISTER_START_ADDRESS].value
        self.bytes_list.append(device_address)
        self.bytes_list.append(self._command.value)
        self.bytes_list.append((register_start_address & 0xff00) >> 8)
        self.bytes_list.append(register_start_address & 0x00ff)

    def _encode_stage_2(self):
        """
        外部勿调用
        编码形成Modbus协议的bytes，存储在self.bytes_list里
        仅实现众山的协议 【众山DTU Modbus协议手册】请求部分
        先执行self._encode_stage_1()获取self._command
        :return: 无返回
        """
        if self._command == ModbusCommandEnum.READ_COIL:
            coil_number = self.decoded[ModbusStructureTypeEnum.COIL_NUMBER]
            self.bytes_list.append((coil_number & 0xff00) >> 8)
            self.bytes_list.append(coil_number & 0x00ff)
        elif self._command == ModbusCommandEnum.READ_DI:
            di_number = self.decoded[ModbusStructureTypeEnum.DI_NUMBER]
            self.bytes_list.append((di_number & 0xff00) >> 8)
            self.bytes_list.append(di_number & 0x00ff)
        elif self._command == ModbusCommandEnum.READ_HOLDING_REGISTER:
            register_number = self.decoded[ModbusStructureTypeEnum.REGISTER_NUMBER]
            self.bytes_list.append((register_number & 0xff00) >> 8)
            self.bytes_list.append(register_number & 0x00ff)
        elif self._command == ModbusCommandEnum.READ_INPUT_REGISTER:
            register_number = self.decoded[ModbusStructureTypeEnum.REGISTER_NUMBER]
            self.bytes_list.append((register_number & 0xff00) >> 8)
            self.bytes_list.append(register_number & 0x00ff)
        elif self._command == ModbusCommandEnum.WRITE_SINGLE_COIL:
            output_value = self.decoded[ModbusStructureTypeEnum.COIL_VALUE]
            if not isinstance(output_value, ModbusWriteSingleCoilValueEnum):
                raise ValueError("Type mismatch, require ModbusWriteSingleCoilValueEnum")
            output_value = output_value.value
            self.bytes_list.append((output_value & 0xff00) >> 8)
            self.bytes_list.append(output_value & 0x00ff)
        elif self._command == ModbusCommandEnum.WRITE_SINGLE_REGISTER:
            register_value = self.decoded[ModbusStructureTypeEnum.REGISTER_VALUE]
            if not isinstance(register_value, ModbusWriteSingleRegisterValueEnum):
                raise ValueError("Type mismatch, require ModbusWriteSingleRegisterValueEnum")
            register_value = register_value.value
            self.bytes_list.append((register_value & 0xff00) >> 8)
            self.bytes_list.append(register_value & 0x00ff)
        elif self._command == ModbusCommandEnum.WRITE_MULTIPLE_COIL:
            coil_number = self.decoded[ModbusStructureTypeEnum.COIL_NUMBER]
            self.bytes_list.append((coil_number & 0xff00) >> 8)
            self.bytes_list.append(coil_number & 0x00ff)
            byte_number = int(ceil(coil_number / 8))
            self.bytes_list.append(byte_number & 0xff)
            output_value = self.decoded[ModbusStructureTypeEnum.COIL_VALUE]
            if not isinstance(output_value, ModbusWriteMultipleCoilValueEnum):
                raise ValueError("Type mismatch, require ModbusWriteMultipleCoilValueEnum")
            output_value = output_value.value
            for i in range(byte_number - 1, -1, -1):
                self.bytes_list.append((output_value >> (i << 3)) & 0xff)
        elif self._command == ModbusCommandEnum.WRITE_MULTIPLE_REGISTER:
            register_number = self.decoded[ModbusStructureTypeEnum.REGISTER_NUMBER]
            self.bytes_list.append((register_number & 0xff00) >> 8)
            self.bytes_list.append(register_number & 0x00ff)
            byte_number = (register_number * 2) & 0xff
            self.bytes_list.append(byte_number)
            register_value = self.decoded[ModbusStructureTypeEnum.REGISTER_VALUE]
            register_value = register_value
            self.bytes_list.extend(register_value)

        bytes_data = bytes(self.bytes_list)
        crc = self.crc16(bytes_data, len(bytes_data))
        self.bytes_list.append((crc & 0xff00) >> 8)
        self.bytes_list.append(crc & 0x00ff)

    def _encode(self):
        """
        外部勿调用
        编码形成Modbus协议的bytes，存储在self.bytes_list和self.blob里
        :return: 无返回
        """
        self._encode_stage_1()
        self._encode_stage_2()
        self.blob = bytes(self.bytes_list)


if __name__ == '__main__':
    # 以下各编码操作对应【众山DTU Modbus协议手册】的示例
    md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: 0x64,
                             ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.READ_COIL,
                             ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZsModbusRegisterAddressEnum.DO1,
                             ModbusStructureTypeEnum.COIL_NUMBER: 0x0004})
    print(md.blob.hex(sep=" "))

    md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: 0x64,
                             ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.READ_DI,
                             ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZsModbusRegisterAddressEnum.DI1,
                             ModbusStructureTypeEnum.DI_NUMBER: 0x0004})
    print(md.blob.hex(sep=" "))

    md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: 0x64,
                             ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.READ_HOLDING_REGISTER,
                             ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZsModbusRegisterAddressEnum.AI1_H,
                             ModbusStructureTypeEnum.REGISTER_NUMBER: 0x0018})
    print(md.blob.hex(sep=" "))

    md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: 0x64,
                             ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.WRITE_SINGLE_COIL,
                             ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZsModbusRegisterAddressEnum.DO2,
                             ModbusStructureTypeEnum.COIL_VALUE: ModbusWriteSingleCoilValueEnum.ON})
    print(md.blob.hex(sep=" "))

    md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: 0x64,
                             ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.WRITE_SINGLE_REGISTER,
                             ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZsModbusRegisterAddressEnum.DO2,
                             ModbusStructureTypeEnum.REGISTER_VALUE: ModbusWriteSingleRegisterValueEnum.ON})
    print(md.blob.hex(sep=" "))

    md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: 0x64,
                             ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.WRITE_MULTIPLE_COIL,
                             ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZsModbusRegisterAddressEnum.DO1,
                             ModbusStructureTypeEnum.COIL_NUMBER: 0x0004,
                             ModbusStructureTypeEnum.COIL_VALUE: ModbusWriteMultipleCoilValueEnum.DO2_ON | ModbusWriteMultipleCoilValueEnum.DO4_ON})
    print(md.blob.hex(sep=" "))

    md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: 0x64,
                             ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.WRITE_MULTIPLE_REGISTER,
                             ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZsModbusRegisterAddressEnum.DO1,
                             ModbusStructureTypeEnum.REGISTER_NUMBER: 0x0004,
                             ModbusStructureTypeEnum.REGISTER_VALUE: [0x00, 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01]})
    print(md.blob.hex(sep=" "))

    # 以下解码对应【众山DTU Modbus协议手册】3.1.1和3.1.2节示例
    md = ModbusData(blob=b"\x64\x01\x01\x0D\x8E\x81")
    print(f"\r\ndevice_address: {hex(md.decoded[ModbusStructureTypeEnum.DEVICE_ADDRESS])}")
    print(f"command: {md.decoded[ModbusStructureTypeEnum.COMMAND].name}")
    print(f"coil0: {md.decoded['coil0'].name}")
    print(f"coil1: {md.decoded['coil1'].name}")
    print(f"coil2: {md.decoded['coil2'].name}")
    print(f"coil3: {md.decoded['coil3'].name}\r\n")

    md = ModbusData(blob=b"\x64\x02\x01\x0C\xBF\x41")
    print(f"\r\ndevice_address: {hex(md.decoded[ModbusStructureTypeEnum.DEVICE_ADDRESS])}")
    print(f"command: {md.decoded[ModbusStructureTypeEnum.COMMAND].name}")
    print(f"di0: {md.decoded['di0'].name}")
    print(f"di1: {md.decoded['di1'].name}")
    print(f"di2: {md.decoded['di2'].name}")
    print(f"di3: {md.decoded['di3'].name}\r\n")

    # 以下为控制串口继电器第一路开
    md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: 0xFE,
                             ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.WRITE_SINGLE_COIL,
                             ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZzModbusRegisterAddressEnum.DO1,
                             ModbusStructureTypeEnum.COIL_VALUE: ModbusWriteSingleCoilValueEnum.ON})
    print(md.blob.hex(sep=" "))

    # 以下为控制串口继电器第十六路关
    md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: 0xFE,
                             ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.WRITE_SINGLE_COIL,
                             ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZzModbusRegisterAddressEnum.DO16,
                             ModbusStructureTypeEnum.COIL_VALUE: ModbusWriteSingleCoilValueEnum.OFF})
    print(md.blob.hex(sep=" "))

    # 以下为读取串口继电器状态
    md = ModbusData(decoded={ModbusStructureTypeEnum.DEVICE_ADDRESS: 0xFE,
                             ModbusStructureTypeEnum.COMMAND: ModbusCommandEnum.READ_COIL,
                             ModbusStructureTypeEnum.REGISTER_START_ADDRESS: ZzModbusRegisterAddressEnum.DO1,
                             ModbusStructureTypeEnum.COIL_NUMBER: 16})
    print(md.blob.hex(sep=" "))

    md = ModbusData(blob=b"\xfe\x01\x02\x00\x00\xad\xe8")
    print(f"\r\ndevice_address: {hex(md.decoded[ModbusStructureTypeEnum.DEVICE_ADDRESS])}")
    print(f"command: {md.decoded[ModbusStructureTypeEnum.COMMAND].name}")
    for i in range(16):
        print(f"coil{i}: {md.decoded['coil' + str(i)].name}")
    print((ModbusReadCoilValueEnum.ON,))
